package com.aispeech.asr.clib;

import com.aispeech.asr.clib.decoder.AudioDecoder;
import com.aispeech.asr.clib.decoder.EAudioCodec;
import com.aispeech.asr.clib.vad.IVadNotifyCallback;
import com.aispeech.asr.clib.vad.Vad;
import com.aispeech.asr.clib.vad.VadSplitter;
import lombok.extern.java.Log;

import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

@Log
public class AudioSplitApp {
    public static void main(String[] args) throws IOException {
        if (args.length != 3) {
            log.warning("args length: " + args.length + ". Usage:");
            log.warning("java -cp <xx>.jar com.aispeech.asr.clib.AudioSplitApp pauseTime resPath <srcAudioFilePath>");
            System.exit(1);
        }

        int pauseTime = Integer.parseInt(args[0]);

        InputStream inputStream = new FileInputStream(args[2]);

        Vad.initialize();

        byte[] data = new byte[1600];
        log.info(String.format("begin split audio: %s, using res: %s ...", args[2], args[1]));


        Vad vad = new Vad(args[1]);
        vad.addCallback((long userData, int state, int offset) -> {
            log.info(String.format("[ Vad.Callback ] offset: %d, state: %d", offset, state));
        });

        int n;
        while ((n = inputStream.read(data)) > 0) {
            vad.feed(data, n);
        }

        vad.stop();
        vad.destroy();
        log.info("Finished.");
    }

    public static void main4(String[] args) throws IOException {
        if (args.length != 4) {
            log.warning("args length: " + args.length + ". Usage:");
            log.warning("java -cp <xx>.jar com.aispeech.asr.clib.AudioSplitApp pauseTime resPath <srcAudioFilePath>");
            System.exit(1);
        }

        int pauseTime = Integer.parseInt(args[0]);
        int maxSpeechTime = 30;

        Path audioPath = Paths.get(args[2]);
        if (!audioPath.toFile().exists()) {
            log.warning("audio not exists. path: " + args[1]);
            System.exit(2);
        }

        byte[] data = Files.readAllBytes(audioPath);
        log.info(String.format("begin split audio: %s ...", args[2]));
        String json = VadSplitter.process(args[1], pauseTime, maxSpeechTime, data, data.length);
        log.info(String.format("[ vad info ]: %s", json));
    }

    public static void main3(String[] args) throws IOException {
        AudioDecoder.initialize();

        if (args.length != 1) {
            log.warning("Usage:");
            log.warning("java -cp <xx>.jar com.aispeech.asr.clib.AudioSplitApp <srcAudioFilePath>");
            System.exit(1);
        }

        Path audioPath = Paths.get(args[0]);
        if (!audioPath.toFile().exists()) {
            log.warning("audio not exists. path: " + args[1]);
            System.exit(2);
        }

        String pcmFileName = "";
        EAudioCodec codec = null;
        int index = args[0].lastIndexOf(".") + 1;

        if (index > 0 && args[0].length() > index) {
            codec = EAudioCodec.from(args[0].substring(index));
            pcmFileName = args[0].substring(0, index - 1) + ".pcm";
        }

        if (codec == null) {
            log.warning(String.format("unknown audio type: %s, path: %s", args[1].substring(index), args[1]));
            System.exit(3);
        }

        byte[] audioBytes = Files.readAllBytes(audioPath);
        AudioDecoder decoder = new AudioDecoder(16000);

        int ret = decoder.start(codec, 16000, 1, 2);
        if (ret != 0) {
            log.warning("decoder start failed. ret: " + ret);
            System.exit(4);
        }

        FileOutputStream outputStream = new FileOutputStream(pcmFileName);
        byte[] decoded = decoder.feed(audioBytes, audioBytes.length);
        outputStream.write(decoded);
        decoded = decoder.stop();
        outputStream.write(decoded);

        decoder.destroy();

        outputStream.close();

        log.info(String.format("Congrats! Audio : %s has decode to PCM: %s", args[0], pcmFileName));
    }

    public static void main2(String[] args) {
        int threads = 2;
        int times = 1;

        if (args.length < 2) {
            log.warning("[ ERROR ] too few arguments.");
            log.warning("Usage:");
            log.warning("java -cp <xx>.jar com.aispeech.asr.clib.AudioSplitApp <vadResPath> <audioFilePath> <nThreads> <nTimes>");
            System.exit(1);
        } else {
            if (args.length >= 3)
                threads = Integer.parseInt(args[2]);

            if (args.length >= 4)
                times = Integer.parseInt(args[3]);
        }

        ExecutorService pool = Executors.newFixedThreadPool(threads);

        AudioDecoder.initialize();
        Vad.initialize();

        Path audioPath = Paths.get(args[1]);
        if (!audioPath.toFile().exists()) {
            log.warning("audio not exists. path: " + args[1]);
            System.exit(2);
        }

        EAudioCodec codec = null;
        int index = args[1].lastIndexOf(".") + 1;

        if (index > 0 && args[1].length() > index) {
            codec = EAudioCodec.from(args[1].substring(index));
        }

        if (codec == null) {
            log.warning(String.format("unknown audio type: %s, path: %s", args[1].substring(index), args[1]));
            System.exit(3);
        }

        log.info(String.format("threads: %d, vad res: %s, audio path: %s", threads, args[0], args[1]));

        try {
            byte[] audioBytes = Files.readAllBytes(audioPath);

            for (int i = 0; i < times; i++) {
                EAudioCodec finalCodec = codec;
                pool.execute(() -> {
                    doAction(audioBytes, args[0], finalCodec);
                });
            }

            pool.shutdown();
            pool.awaitTermination(8, TimeUnit.HOURS);
        } catch (IOException e) {
            e.printStackTrace();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    private static void doAction(byte[] audioBytes, String vadResPath, EAudioCodec codec) {
        AudioDecoder decoder;
        Vad vad;
        try {
            decoder = new AudioDecoder(16000);
            vad = new Vad(vadResPath);

            List<int[]> states = new ArrayList<>();

            vad.addCallback((userData, state, offset) -> {
                states.add(new int[]{offset, state});
//                log.info(String.format("[ callback ] offset: %d, state: %d", state, offset));
            });

            int ret = -1;

            ret = decoder.start(codec, 16000, 1, 2);
            if (ret != 0) {
                log.warning("decoder start failed. ret: " + ret);
                return;
            }

            ret = vad.start();
            if (ret != 0) {
                log.warning("vad start failed. ret: " + ret);
                return;
            }

            int total = audioBytes.length;
            int offset = 0;
            byte[] decoded;

            while (offset < total) {
                int len = Math.min(3200, total - offset);
                byte[] bytes = new byte[len];

                for (int i = 0; i < len; i++) bytes[i] = audioBytes[offset + i];

                decoded = decoder.feed(bytes, len);
                vad.feed(decoded, decoded.length);

                offset += len;
            }

            decoded = decoder.stop();

            vad.feed(decoded, decoded.length);
            vad.stop();
            vad.destroy();

            decoder.destroy();

            log.info(String.format("states count: %d", states.size()));




        } catch (RuntimeException e) {
            e.printStackTrace();
        }
    }
}
